# Copyright (c) 2022 Presto Labs Pte. Ltd.
# Author: lucasyoo
"""
Related to OTC Option Pricing
"""
import datetime
import pandas as pd
import numpy as np
from scipy.stats import norm
import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import json
import subprocess
import os
import git
from collections import (defaultdict, OrderedDict)

import coin.pnl.sim_stat_plot as ssplot
from coin.base.datetime_util import (to_datetime, to_timestamp_int)
import coin.pnl.sim_stat_frompb as ssfrompb


def format_table(table_df):
  big_num = {'total_value(USD)', 'trade_pnl(USD)', 'option_value(USD)', 'call_num'}
  percent = {'risk_free_rate', 'volatility'}
  original_column = list(table_df.columns)
  tmp_column = original_column.copy()
  for idx, col in enumerate(original_column):
    if col in big_num:
      table_df[col] = table_df[col].apply(lambda x: f'{x:,.4f}')
    if col in percent:
      table_df[col] = table_df[col].apply(lambda x: f'{x:.2%}')
      table_df.rename(columns={col: col + '(%)'}, inplace=True)
      tmp_column[idx] = col + '(%)'

  return table_df[tmp_column]


def plot_option_and_hedge_pnl(dir, name, total_values, option_values, trade_pnls, mtm_pos, deltas,
                              ttms, table_config_df, table_stat_df, symbol):
  plt.rcParams["figure.figsize"] = 15, 17
  ssplot.setup_plt()
  plt.subplot2grid((10, 1), (0, 0), rowspan=2)
  color_list = ['r', 'b', 'g', 'c', 'm', 'y']
  pnl_df = None
  mtms = {k: v["mtm_price"] for k, v in mtm_pos.items()}
  if symbol is None:
    if len(total_values) <= 1:
      return
    total_pnl = pd.concat(list(total_values.values()),
                          axis=1).sort_index().ffill().bfill().sum(axis=1)
    total_option_value = pd.concat(list(option_values.values()),
                                   axis=1).sort_index().ffill().bfill().sum(axis=1)
    total_trade_pnl = pd.concat(list(trade_pnls.values()), axis=1).sort_index().ffill().sum(axis=1)
    pnl_df = pd.concat([total_trade_pnl, total_option_value, total_pnl],
                       axis=1,
                       keys=['trade_pnl', 'option_value',
                             'total_value']).sort_index().ffill().bfill()
    plt.plot(pd.to_datetime(total_pnl.index),
             total_pnl.to_numpy(),
             'k-',
             label='total_value(USD)',
             lw=0.5,
             drawstyle='steps-post')
    plt.plot(pd.to_datetime(total_option_value.index),
             total_option_value.to_numpy(),
             'r-',
             label='option_value(USD)',
             lw=0.2,
             drawstyle='steps-post')
    plt.plot(pd.to_datetime(total_trade_pnl.index),
             total_trade_pnl.to_numpy(),
             'b-',
             label='trade_pnl(USD)',
             lw=0.2,
             drawstyle='steps-post')
    plt.legend(loc='upper right')

    plt.subplot2grid((10, 1), (2, 0), rowspan=2)
    plt.plot(pd.to_datetime(total_pnl.index),
             total_pnl.to_numpy(),
             'k-',
             label='Agg_total_value(USD)',
             lw=0.5,
             drawstyle='steps-post')
    for i, (k, v) in enumerate(total_values.items()):
      # sliced_idx = v.index.intersection(ts_idx)
      plt.plot(pd.to_datetime(v.index),
               v.to_numpy(),
               color_list[i],
               label=f'{k}_total_value(USD)',
               lw=0.2,
               drawstyle='steps-post')
    plt.legend(loc='upper right')
    plt.twinx()
    for i, (_, v) in enumerate(mtms.items()):
      plt.plot(pd.to_datetime(v.index), v.to_numpy(), color_list[i], lw=0.2, alpha=0.5)
    plt.subplot2grid((10, 1), (4, 0), rowspan=2)
    plt.plot(pd.to_datetime(deltas['Agg'].index),
             deltas['Agg'].to_numpy(),
             'k-',
             label='Agg_delta',
             lw=0.5,
             drawstyle='steps-post')
    for i, (k, delta) in enumerate(deltas.items()):
      if k == 'Agg':
        continue
      plt.plot(pd.to_datetime(delta.index),
               delta.to_numpy(),
               color_list[i],
               label=f'{k}_delta',
               lw=0.2,
               drawstyle='steps-post')
    plt.legend(loc='upper right')
    plt.twinx()
    for i, (k, _) in enumerate(deltas.items()):
      if k == 'Agg':
        continue
      plt.plot(pd.to_datetime(mtms[k].index), mtms[k].to_numpy(), color_list[i], lw=0.2, alpha=0.5)
    plt.subplot2grid((10, 1), (6, 0), rowspan=2)
    idx = 0
    for k, symbol_ttm in ttms.items():
      for option_name, ttm in symbol_ttm.items():
        plt.plot(pd.to_datetime(ttm.index),
                 ttm.to_numpy(),
                 color_list[idx],
                 label=f'{k}_{option_name}_ttm',
                 lw=0.5,
                 drawstyle='steps-post')
        idx += 1
    plt.legend(loc='upper right')
  else:
    pnl_df = pd.concat([trade_pnls[symbol], option_values[symbol], total_values[symbol]],
                       axis=1,
                       keys=['trade_pnl', 'option_value',
                             'total_value']).sort_index().ffill().bfill()
    plt.plot(pd.to_datetime(total_values[symbol].index),
             total_values[symbol].to_numpy(),
             'k-',
             label='total_value(USD)',
             lw=0.5,
             drawstyle='steps-post')
    plt.plot(pd.to_datetime(option_values[symbol].index),
             option_values[symbol].to_numpy(),
             'r-',
             label='option_value(USD)',
             lw=0.2,
             drawstyle='steps-post')
    plt.plot(pd.to_datetime(trade_pnls[symbol].index),
             trade_pnls[symbol].to_numpy(),
             'b-',
             label='trade_pnl(USD)',
             lw=0.2,
             drawstyle='steps-post')
    plt.legend(loc='upper right')
    plt.twinx()
    plt.plot(pd.to_datetime(mtms[symbol].index), mtms[symbol].to_numpy(), 'c-', lw=0.2, alpha=0.5)
    plt.subplot2grid((10, 1), (2, 0), rowspan=2)
    plt.plot(pd.to_datetime(total_values[symbol].index),
             total_values[symbol].to_numpy(),
             'k-',
             label='total_value(USD)',
             lw=0.5,
             drawstyle='steps-post')
    plt.legend(loc='upper right')
    plt.twinx()
    plt.plot(pd.to_datetime(mtms[symbol].index),
             mtms[symbol].to_numpy(),
             'c-',
             lw=0.2,
             alpha=0.5,
             drawstyle='steps-post')
    plt.subplot2grid((10, 1), (4, 0), rowspan=2)
    plt.plot(pd.to_datetime(deltas[symbol].index),
             deltas[symbol].to_numpy(),
             'k-',
             label=f'{symbol}_delta',
             lw=0.5,
             drawstyle='steps-post')
    plt.legend(loc='upper right')
    plt.twinx()
    plt.plot(pd.to_datetime(mtms[symbol].index),
             mtms[symbol].to_numpy(),
             'c-',
             lw=0.2,
             alpha=0.5,
             drawstyle='steps-post')
    plt.subplot2grid((10, 1), (6, 0), rowspan=2)
    for idx, (option_name, ttm) in enumerate(ttms[symbol].items()):
      plt.plot(pd.to_datetime(ttm.index),
               ttm.to_numpy(),
               color_list[idx],
               label=f'{option_name}_ttm',
               lw=0.5,
               drawstyle='steps-post')
    plt.legend(loc='upper right')

  for idx, df in enumerate([table_config_df, table_stat_df]):
    ax = plt.subplot2grid((10, 1), (8 + idx, 0), rowspan=1)
    ax.axis('tight')
    ax.axis('off')
    formated_table = format_table(df)
    table = ax.table(cellText=formated_table.values,
                     colLabels=formated_table.columns,
                     loc='center',
                     cellLoc='center')
    table.auto_set_column_width(col=list(range(len(formated_table.columns))))
  plt.suptitle(
      f'{name}\n'
      f'symbols: {symbol if symbol is not None else ", ".join(list(total_values.keys()))}\n'
      f'{pd.to_datetime(pnl_df.index).min()}-{pd.to_datetime(pnl_df.index).max()}',
      fontsize=12)
  plt.savefig(f'{dir}/{symbol if symbol is not None else "total"}_{name}.png')
  plt.close()
  pnl_df.to_json(f"{dir}/{symbol if symbol is not None else 'total'}_{name}.json", orient="split")
  pd.concat([df.assign(symbol=symbol) for symbol, df in mtm_pos.items()
            ]).to_json(f"{dir}/mtm_pos.json", orient="split")


def plot_options_and_hedge_pnls(sim_result_dir, pnldf, sim_config_path, is_pnlplot=False):
  """
  Plot Option and Hedge PnL With 1min interval
  Assume (1) Spot  (2) USDT or USD quote_currency
  """
  start_ts = int(pnldf['timestamp'].min())
  end_ts = int(pnldf['timestamp'].max())
  opt_configs = {}
  if is_pnlplot:
    subprocess.run(['git', 'switch', 'master'], cwd=os.path.expandvars('$COIN_DEPLOY_REPO'))
    subprocess.run(['git', 'fetch', 'origin', 'master'],
                   cwd=os.path.expandvars('$COIN_DEPLOY_REPO'))
    repo = git.Repo(os.path.expandvars("$COIN_DEPLOY_REPO"))
    strategy = sim_config_path.replace(".json", "")
    strat_group = strategy.split("_")[0]
    if strat_group == "vmm" and "delta" in strategy:
      commits = list(repo.iter_commits(paths=f"{strat_group}/config/{strategy}.json"))
      for commit in commits:
        utc_commit_time = commit.committed_datetime.astimezone(
            datetime.timezone.utc).replace(tzinfo=None)
        if utc_commit_time > to_datetime(end_ts):
          continue
        for blob in commit.tree.traverse():
          if blob.path == f"{strat_group}/config/{strategy}.json":
            try:
              strat_config = json.loads(blob.data_stream.read().decode("utf-8"))
            except json.JSONDecodeError:
              continue
            opt_configs[to_timestamp_int(utc_commit_time)] = strat_config
        if utc_commit_time < to_datetime(start_ts):
          break
  else:
    config_path = f"{sim_result_dir}/configs/{sim_config_path}" if not is_pnlplot else os.path.join(
        os.path.expandvars("$COIN_REPO"), 'python', 'coin_deploy',
        sim_config_path.split("_")[0], 'config', sim_config_path)
    with open(config_path, "r") as f:
      sim_config = json.load(f)
      opt_configs[start_ts] = sim_config
  option_dir = os.path.join(sim_result_dir, "option_plot",
                            f"{sim_config_path.replace('.json', '')}")
  os.makedirs(option_dir, exist_ok=True)
  if not opt_configs:
    total_pnl = pnldf.groupby(["timestamp"])["pnl_diff(USD)"].sum().cumsum().ffill().dropna()
    total_pnl.name = "trade_pnl"
    total_pnl = total_pnl.to_frame()
    total_pnl.index.name = None
    total_pnl["total_pnl"] = total_pnl["trade_pnl"]
    total_pnl.to_json(f"{option_dir}/total_pnl.json", orient="split")
    for symbol in pnldf['from_symbol'].unique():
      tmpdf = pnldf[pnldf['from_symbol'] == symbol].copy()
      tmp_pnl = tmpdf.groupby(["timestamp"])["pnl_diff(USD)"].sum().cumsum().ffill().dropna()
      tmp_pnl.name = "trade_pnl"
      tmp_pnl = tmp_pnl.to_frame()
      tmp_pnl.index.name = None
      tmp_pnl["total_pnl"] = tmp_pnl["trade_pnl"]
      tmp_pnl.to_json(f"{option_dir}/{symbol}_pnl.json", orient="split")
    return
  option_infos = defaultdict(dict)
  for dt, sim_config in opt_configs.items():

    if 'oc_config' not in sim_config or (
        'oc_config' in sim_config and 'configs' in sim_config['oc_config'] and not [
            i.keys() for i in list(sim_config['oc_config']['configs'].values()) if 'contracts' in i
        ]):
      continue
    symbol_info = sim_config['oc_config']['configs']
    for symbol, info in symbol_info.items():
      if 'contracts' in info:  # multiple expiry
        option_infos[symbol][dt] = info['contracts']
      else:
        option_infos[symbol][dt] = {"1": {"expiry": info['expiry'], "call_num": info['call_num']}}
        if "strike_price_num" in info:
          option_infos[symbol][dt]['strike_price_num'] = info[
              'report_strike_price_num'] if 'report_strike_price_num' in info else info[
                  'strike_price_num']
          option_infos[symbol][dt]['strike_price_denom'] = info[
              'report_strike_price_denom'] if 'report_strike_price_denom' in info else info[
                  'strike_price_denom']
        else:
          option_infos[symbol][dt]["strike_price"] = info["strike_price"]
      option_infos[symbol][dt]["risk_free_rate"] = info['risk_free_rate']
      option_infos[symbol][dt]["volatility"] = info['volatility']
      option_infos[symbol][dt]["bandwidth_bp"] = info['bandwidth_bp']
      option_infos[symbol][dt]["layer_num"] = info['layer_num']

  total_values = {}
  total_option_values = {}
  total_trade_pnls = {}
  total_mtm_pos = {}
  total_ttms = defaultdict(dict)

  num = []
  denom = []
  agg_deltas = {}

  table_config_datas = defaultdict(OrderedDict)
  table_stat_datas = []
  for symbol, symbol_info in option_infos.items():
    from_symbol = '.'.join(symbol.split(":")[1:])
    target_df = pnldf[pnldf['from_symbol'] == from_symbol].sort_values('timestamp')
    target_df['pnl_net(USD)'] = target_df['pnl_diff(USD)'].cumsum()
    target_df.reset_index(drop=True, inplace=True)
    target_df['mtm_price'].ffill(inplace=True)
    target_df['mtm_price'].bfill(inplace=True)
    target_df['open_pos'].ffill(inplace=True)
    target_df['open_pos'].bfill(inplace=True)
    option_values = []
    opt_dts = sorted(list(symbol_info.keys()))
    delta_info = defaultdict(list)
    ttm_info = defaultdict(list)

    for tmp_idx in range(len(opt_dts)):
      opt_dt = opt_dts[tmp_idx]
      next_dt = np.inf if tmp_idx == len(opt_dts) - 1 else opt_dts[tmp_idx + 1]
      tmp_target_df = target_df[(target_df["timestamp"] >= opt_dt) &
                                (target_df["timestamp"] < next_dt)].copy()
      if tmp_target_df.empty:
        continue
      info = symbol_info[opt_dt]
      risk_free_rate = info['risk_free_rate']
      volatility = info['volatility']
      layer_bandwidth_bp = info['bandwidth_bp']
      layer_num = info['layer_num']
      total_mtm_pos[symbol] = target_df[['timestamp', 'mtm_price',
                                         'open_pos']].ffill().dropna().set_index('timestamp')

      tmp_option_values = []
      for idx, option_info in info.items():
        if idx in {'risk_free_rate', 'volatility', 'bandwidth_bp', 'layer_num', 'owner'}:
          continue
        expiry = option_info['expiry']
        expiry_ts = pd.Timestamp(expiry).value
        call_num = option_info['call_num']
        if 'report_strike_price_num' in option_info:
          strike_price = float(option_info['report_strike_price_num']) / float(
              option_info['report_strike_price_denom'])
        elif "strike_price_num" in option_info:
          strike_price = float(option_info['strike_price_num']) / float(
              option_info['strike_price_denom'])
        else:
          strike_price = float(option_info['strike_price'])

        ttm = (expiry_ts - tmp_target_df['timestamp']) / 1e9 / 3600 / 24 / 365
        d1 = (np.log(tmp_target_df['mtm_price'] / strike_price) +
              (risk_free_rate + 0.5 * volatility**2) * ttm) / (volatility * np.sqrt(ttm))
        if d1.dropna().empty:
          continue
        d2 = d1 - volatility * np.sqrt(ttm)
        call_price = tmp_target_df['mtm_price'] * norm.cdf(d1) - strike_price * np.exp(
            -risk_free_rate * ttm) * norm.cdf(d2)
        option_value = (call_num * call_price)
        option_value.index = tmp_target_df['timestamp']
        option_value.name = idx
        option_value = option_value[option_value.index <= expiry_ts].ffill().dropna()
        tmp_option_values.append(option_value)
        delta = call_num * pd.Series(norm.cdf(d1), index=option_value.index).to_frame()
        delta['call_num'] = call_num
        delta_info[f'{idx}_T_{expiry}_K_{strike_price}'].append(delta)
        ttm.index = option_value.index
        ttm_info[f'{idx}_T_{expiry}_K_{strike_price}'].append(ttm)
        table_config_datas[f'{symbol}_{idx}_T_{expiry}_K_{strike_price}'] = OrderedDict(
            (('expiry', expiry), ('strike_price', strike_price), ('call_num', call_num),
             ('risk_free_rate', risk_free_rate), ('volatility', volatility),
             ('layer_bandwidth_bp', layer_bandwidth_bp), ('layer_num', layer_num)))
      option_values.append(pd.concat(tmp_option_values, axis=1).sort_index())
    symbol_denom = []
    symbol_num = []
    for _, vals in delta_info.items():
      tmp_df = pd.concat(vals)
      tmp_df = tmp_df.loc[~tmp_df.index.duplicated(keep='first')].sort_index().ffill()
      symbol_denom.append(tmp_df['call_num'])
      symbol_num.append(tmp_df[[i for i in tmp_df.columns if i != 'call_num'][0]] *
                        tmp_df['call_num'])
    tmp_nom = pd.concat(symbol_num, axis=1).sort_index().ffill().dropna().sum(axis=1)
    tmp_denom = pd.concat(symbol_denom, axis=1).sort_index().ffill().dropna().sum(axis=1)
    tmp_nom = tmp_nom[~tmp_nom.index.duplicated(keep='first')]
    tmp_denom = tmp_denom[~tmp_denom.index.duplicated(keep='first')]
    intersec_idx = tmp_nom.index.intersection(tmp_denom.index)
    tmp_nom = tmp_nom.loc[intersec_idx]
    tmp_denom = tmp_denom.loc[intersec_idx]
    agg_deltas[symbol] = tmp_nom / tmp_denom
    denom.append(tmp_denom)
    num.append(tmp_nom)
    for k, v in ttm_info.items():
      tmp_df = pd.concat(v).sort_index().ffill()
      total_ttms[symbol][k] = tmp_df
    option_value_df = pd.concat(option_values).sort_index().ffill()
    if option_value_df.empty:
      continue
    total_option_value = option_value_df.sum(axis=1)
    total_option_value.name = 'option_value(USD)'
    trade_pnl = pd.Series(target_df['pnl_net(USD)'].values,
                          index=target_df['timestamp'],
                          name='trade_pnl(USD)')
    trade_pnl = trade_pnl[~trade_pnl.index.duplicated(keep='first')]
    total_option_value = total_option_value[~total_option_value.index.duplicated(keep='first')]
    quote = symbol.split(":")[-1].split(".")[0].split("-")[1]

    dts = np.sort(
        np.unique(
            pd.to_datetime(list(set(total_option_value.index).union(set(
                trade_pnl.index)))).date)).tolist()
    fxrate_sr = ssfrompb.get_fxrate(quote, tuple(dts))
    total_option_value *= (1.0 / pd.Series(pd.to_datetime(
        total_option_value.index).date).map(fxrate_sr)).to_numpy()
    trade_pnl *= (1.0 / pd.Series(pd.to_datetime(trade_pnl.index).date).map(fxrate_sr)).to_numpy()
    tmpdf = pd.concat([total_option_value, trade_pnl], axis=1).sort_index().ffill()
    tmpdf["trade_pnl(USD)"].fillna({tmpdf.index[0]: 0.0}, inplace=True)
    tmpdf["trade_pnl(USD)"].ffill(inplace=True)
    tmpdf['option_value(USD)'].ffill(inplace=True)
    tmpdf["option_value(USD)"].bfill(inplace=True)
    tmpdf["option_value(USD)"].fillna(0.0, inplace=True)
    trade_pnl = tmpdf['trade_pnl(USD)']
    option_value = tmpdf['option_value(USD)']
    total_value = tmpdf.fillna(0.0).sum(axis=1)
    total_value.name = 'total_value(USD)'
    total_values[symbol] = total_value
    total_option_values[symbol] = option_value
    total_trade_pnls[symbol] = trade_pnl
    stat_lst = []
    for idx, date in enumerate(['start', 'end']):
      stat_lst.append(
          OrderedDict((('date', date), ('total_value(USD)', total_value.iloc[-idx]),
                       ('trade_pnl(USD)', trade_pnl.iloc[-idx]), ('option_value(USD)',
                                                                  option_value.iloc[-idx]))))
    config_df = pd.DataFrame.from_dict(table_config_datas, orient='index')
    config_df.index.name = 'option_name'
    stat_df = pd.DataFrame(stat_lst)
    diff = stat_df[['total_value(USD)', 'trade_pnl(USD)', 'option_value(USD)']].diff().iloc[-1]
    diff.loc['date'] = "diff"
    stat_df = pd.concat([stat_df, diff.to_frame().T], sort=False)
    plot_option_and_hedge_pnl(option_dir, sim_config_path.replace('.json', ''), total_values,
                              total_option_values, total_trade_pnls, total_mtm_pos, agg_deltas,
                              total_ttms,
                              config_df[config_df.index.str.startswith(symbol)].reset_index(),
                              stat_df, symbol)
    table_stat_datas.extend(stat_lst)
  denom = pd.concat(denom, axis=1).sort_index().ffill().dropna().sum(axis=1)
  num = pd.concat(num, axis=1).sort_index().ffill().dropna().sum(axis=1)
  intersec_idx = denom.index.intersection(num.index)
  denom = denom.loc[intersec_idx]
  num = num.loc[intersec_idx]
  agg_deltas['Agg'] = num / denom
  config_df = pd.DataFrame(table_config_datas).T
  config_df.index.name = 'option_name'
  stat_df = pd.DataFrame(table_stat_datas)
  stat_df = stat_df.groupby('date', as_index=False).sum().sort_values('date', ascending=False)
  diff = stat_df[['total_value(USD)', 'trade_pnl(USD)', 'option_value(USD)']].diff().iloc[-1]
  diff.loc['date'] = "diff"
  stat_df = pd.concat([stat_df, diff.to_frame().T], sort=False)
  plot_option_and_hedge_pnl(option_dir, sim_config_path.replace('.json', ''), total_values,
                            total_option_values, total_trade_pnls, total_mtm_pos, agg_deltas,
                            total_ttms, config_df.reset_index(), stat_df, None)
